#include "trace.h"
#include "slep.h"


// A is p*n; B is p*m; W is n*m; p>=m
// MINIMIZE f(W)= |A*W - B|_F^2 + |W|_*
TRACERET mat_primal(double* A, double* B, double lambda, TRACEOPTS opt){
	long maxIter,iter,i,j;
	double meps;
	int m,n,p, rank, info, rm1;
	int reduce_flag;
	int r, freq;

	double *sing;
	double *q0, *r0;

	double *M, *M2, m3;
	double *Anew, *Bnew, *C, *E, *W, *W0, *wr, *wi;
	double *Y, *G, *T, *R, *ST, *D, *TMP1 , *U;
	double f, theta, theta0, L, normWW0;
	double pobj, dobj, tol;
	TRACERET ret;

	m = opt.m;
	n = opt.n;
	p = opt.p;

	//E = (double *) malloc( sizeof(double) * );
	//W = (double *) malloc( sizeof(double) * );
	//W0 = (double *) malloc( sizeof(double) * );
	//wi = (double *) malloc( sizeof(double) * );
	//wr = (double *) malloc( sizeof(double) * );
	//Y = (double *) malloc( sizeof(double) * );
	//G = (double *) malloc( sizeof(double) * );
	//T = (double *) malloc( sizeof(double) * );
	//U = (double *) malloc( sizeof(double) * );


	dscal(p*n,1/sqrt(lambda),A,1);
	dscal(p*m,1/sqrt(lambda),B,1);
	
	maxIter = 1000000;
	meps = 1e-10;
	iter=0;

	reduce_flag=1;
	if(p>=n){
		rank= rankofMat(A,p,n);
		if(rank==n)
			reduce_flag=0;
	}

	
	if(reduce_flag==1){
		// reduce A to have full row rank
		q0 = (double *) malloc( sizeof(double) * n*n);
		r0 = (double *) malloc( sizeof(double) * n*p);
		slepQRFactorize('T',p,n,A,q0,r0);
		r = rankofMat(r0,n,p);
		Anew = (double *) malloc( sizeof(double) * p * r );
		Bnew = (double *) malloc( sizeof(double) * p * m );
		dcopyMatrixTranspose(r,p,r0,n,Anew,p);
		dcopy(p*m,B,1,Bnew,1);
	}else
		r=n;
	
	if( reduce_flag == 1){
		M = (double *) malloc( sizeof(double) * r * r);
		M2 = (double *) malloc( sizeof(double) * r * m);
		dgemm('T','N',r,r,p,1.0,Anew,p,Anew,p,0,M,r);
		dgemm('T','N',r,m,p,1.0,Anew,p,Bnew,p,0,M2,r);
		m3 = ddot(p*m,Bnew,1,Bnew,1);
		free(Anew);
		free(Bnew);
	}else{
		M = (double *) malloc( sizeof(double) * r * r);
		M2 = (double *) malloc( sizeof(double) * r * m);
		dgemm('T','N',r,r,p,1.0,A,p,A,p,0,M,r);
		dgemm('T','N',r,m,p,1.0,A,p,B,p,0,M2,r);
		m3 = ddot(p*m,B,1,B,1);
	}
	C = (double *) malloc( sizeof(double) * r*r);
	dcopy(r*r,M,1,C,1);
	slepInvert(r,C);
	E = (double *) malloc( sizeof(double) * m*r);
	dgemm('T','N',m,r,r,1.0,M2,r,C,r,0,E,m);
	f=0;
	for(i=0;i<m;i++)
		f += ddot(r,&M2[i*r],1,&E[i],m);
	f = f - m3;

	dcopyMatrixTranspose(m,r,E,m,W,r);

	dcopy(r*m,W,1,W0,1);

	theta=1;
	theta0=1;

	dgeev('N','N',r,M,r,wr,wi,NULL,1,NULL,1,&info);
	L = wr[largestMagnitude(r,wr,wi)];

	for(iter=0;iter<maxIter; iter++){
		dcopy(r*m,W,1,Y,1);
		dscal(r*m,1 + theta*(1/theta0 -1),Y,1);
		daxpy(r*m,-theta*(1/theta0 -1),W0,1,Y,1);

		// G= M*Y - M2;
		dcopy(r*m,M2,1,G,1);
		dgemm('N','N',r,m,r,1,M,r,Y,r,-1,G,r);

		//T=Y-G/L;
		dcopy(r*m,Y,1,T,1);
		daxpy(r*m,-1/L,G,1,T,1);
		
		rm1 = minof2(r,m);
		D = (double *) malloc(sizeof(double) * rm1);
		R = (double *) malloc(sizeof(double) * r * rm1);
		ST = (double *) malloc(sizeof(double) * rm1 * m);
		slepSVD(r,m,T,r,D,R,r,ST,rm1);
		//W0=W;
		dcopy(r*m,W,1,W0,1);

		//W=R*(max(D-eye(size(D,1))/L,0))*S';
		daxpy(rm1,-1/L,oneVector,1,D,1);
		TMP1 = (double *) malloc(sizeof(double) * r*rm1);
		dcopy(r*rm1,R,1,TMP1,1);
		for(i=0;i<rm1;i++){
			D[i] = maxof2(0,D[i]);
			dscal(r,D[i],&TMP1[i*r],1);
		}

		dgemm('N','N',r,m,rm1,1,TMP1,r,ST,rm1,0,W,r);
		free(TMP1);

		normWW0 = sqrt( ddot(r*m,W,1,W,1)
			+ddot(r*m,W,1,W0,1) - 2*ddot(r*m,W,1,W0,1) );
		if( normWW0 <= meps){
			//printf("termination due to negligible change in U= %lf",normWW0);
			//U= (W'-E) * M;
			TMP1 = (double *) malloc(sizeof(double) * r*m);
			dcopyMatrixTranspose(r,m,W,r,TMP1,m);
			daxpy(m*r,-1,E,m,TMP1,m);
			dgemm('N','N',m,r,r,1,TMP1,m,M,r,0,U,m);
			//[R,D,S]=svd(U,'econ');
			/* we are going to re-use R and ST, as they have 
			the same memory allocation but swapped*/
			R  = ((long)R  ^ (long)ST) ;
			ST = ((long)ST ^ (long)R ) ;
			R  = ((long)R  ^ (long)ST );
			slepSVD(m,r,U,m,D,R,m,ST,rm1);
			TMP1 = (double *) malloc(sizeof(double) * m*rm1);
			dcopy(m*rm1,R,1,TMP1,1);
			for(i=0;i<rm1;i++){
				D[i] = minof2(1.0l,D[i]);
				dscal(m,D[i],&TMP1[i*m],1);
			}
			dgemm('N','N',m,r,rm1,1,TMP1,m,ST,rm1,0,U,m);
			free(TMP1);


			pobj = slepNuclearNorm(r,m,W,r) + (myTrace3('T',m,r,r,W,r,M,r,W,r)
				-2*myTrace2('T','N',m,r,M2,r,W,r)+m3)/2;

			dobj = -(myTrace2('T','N',r,m,E,m,U,m) + (myTrace3('T',r,m,r,U,m,U,m,C,r)+f)/2);
		}
		if(iter>0 && (iter%freq==0)){
			//printf("termination due to negligible change in U= %lf",normWW0);
			//U= (W'-E) * M;
			TMP1 = (double *) malloc(sizeof(double) * r*m);
			dcopyMatrixTranspose(r,m,W,r,TMP1,m);
			daxpy(m*r,-1,E,m,TMP1,m);
			dgemm('N','N',m,r,r,1,TMP1,m,M,r,0,U,m);
			//[R,D,S]=svd(U,'econ');
			/* we are going to re-use R and ST, as they have 
			the same memory allocation but swapped*/
			R  = (long)R  ^ (long)ST ;
			ST = (long)ST ^ (long)R  ;
			R  = (long)R  ^ (long)ST ;
			slepSVD(m,r,U,m,D,R,m,ST,rm1);
			TMP1 = (double *) malloc(sizeof(double) * m*rm1);
			dcopy(m*rm1,R,1,TMP1,1);
			for(i=0;i<rm1;i++){
				D[i] = minof2(1.0l,D[i]);
				dscal(m,D[i],&TMP1[i*m],1);
			}
			dgemm('N','N',m,r,rm1,1,TMP1,m,ST,rm1,0,U,m);
			free(TMP1);


			pobj = slepNuclearNorm(r,m,W,r) + (myTrace3('T',m,r,r,W,r,M,r,W,r)
				-2*myTrace2('T','N',m,r,M2,r,W,r)+m3)/2;

			dobj = -(myTrace2('T','N',r,m,E,m,U,m) + (myTrace3('T',r,m,r,U,m,U,m,C,r)+f)/2);
			if( fabs(pobj-dobj) < tol * ( fabs(dobj) + 1) )
				break;
		}
	}
	
	ret.m = r;
	ret.n = m;
	if( reduce_flag==1){
		//W=R0*[W;zeros(n-r,m)];
		TMP1 = (double*) malloc(sizeof(double) *n *m); 
		memset(TMP1,0, n*m * sizeof(double) );
		//dcopy(r*m,W,1,TMP1,1);
		dgemm('N','N',n,m,r,1,q0,n,W,r,0,TMP1,n);
		free(W);
		W=TMP1;
		TMP1=NULL;
		ret.m = n;
		ret.n = m;
	}
	ret.W = W;
	ret.pobj = pobj;
	return ret;
}